import dgl
import torch
import numpy as np
import json
import torch.nn.functional as F
import dgl.data
from build_unknow_malware_graph import build_graph_func
from collections import Counter
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score


# from model import SAGE
# from model import GAT
from model import GCN

def main():
    family_file_path = './syscall_family.json'
    unknown_family_str = "trojan.linux/gafgyt"

    g,sorted_family_keys = build_graph_func(data_file_path='./combine.json',family_file_path=family_file_path,unknown_family_str=unknown_family_str)
    g = g.to('cuda')
    g=dgl.add_self_loop(g)

    # Get the labels
    labels = g.ndata['label']
    benign_samples = torch.where(labels == 0)[0]
    malicious_samples = torch.where(labels == 1)[0]

    # Initialize counters
    benign_syscall_count = 0
    malicious_syscall_count = 0

    # Count the number of system calls for benign samples
    for sample in benign_samples:
        # Count the number of edges (system calls) for each benign sample
        benign_syscall_count += g.in_degrees(sample).item()

    # Count the number of system calls for malicious samples
    for sample in malicious_samples:
        # Count the number of edges (system calls) for each malicious sample
        malicious_syscall_count += g.in_degrees(sample).item()


    in_feats = g.ndata["feat"].shape[1]
    num_classes = 2

    model = GCN(in_feats, num_classes).to('cuda')

    unknown_family_id = sorted_family_keys.index(unknown_family_str)
    train_unknown(g, model,unknown_family_id)

    # Calculate averages
    # avg_benign_syscalls = benign_syscall_count / len(benign_samples) if len(benign_samples) > 0 else 0
    # avg_malicious_syscalls = malicious_syscall_count / len(malicious_samples) if len(malicious_samples) > 0 else 0

    # print(f"Average number of system calls for benign samples: {avg_benign_syscalls:.2f}")
    # print(f"Average number of system calls for malicious samples: {avg_malicious_syscalls:.2f}")
    # exit()
    # family_count = Counter(family_file.values())
    # sorted_family_count = sorted(family_count.items(), key=lambda x: x[1], reverse=True)
    # sorted_family_keys = [item[0] for item in sorted_family_count] # 0:136 共137个家族

    # print("Node feature")
    # print(g.ndata)

    # print("Edge feature")
    # print(g.edata)

def train_unknown(g, model,unknown_family_id, num_epoch = 100, learning_rate =  0.001):
    """
    train function
    """
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    best_test_accurate = 0
    
    features = g.ndata["feat"]
    labels = g.ndata["label"].long()
    train_mask = g.ndata["train_mask"] # 1264个true
    test_mask = g.ndata["test_mask"] # 317个true
    # print(train_mask.sum().item())
    # exit()


    for e in range(num_epoch):
        #forward
        result = model(g, features)
        #prediction
        pred = result.argmax(dim=1)
        #Loss
        loss = F.cross_entropy(result[train_mask], labels[train_mask])
        #compute accurate
        train_accurate = (pred[train_mask]==labels[train_mask]).float().mean()
        test_accurate = (pred[test_mask]==labels[test_mask]).float().mean()
        if best_test_accurate < test_accurate:
            best_test_accurate = test_accurate
        #backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if e % 1 == 0:
            misclassified = pred[test_mask] != labels[test_mask]
            misclassified_labels = labels[test_mask][misclassified]
            misclassified_count = Counter(misclassified_labels.cpu().numpy())
            rightclassified = pred[test_mask] == labels[test_mask]
            rightclassified_labels = labels[test_mask][rightclassified]
            rightclassified_count = Counter(rightclassified_labels.cpu().numpy())

            benign_misclassified_count = misclassified_count.get(0, 0)
            malware_misclassified_count = misclassified_count.get(1, 0)
            benign_rightclassified_count = rightclassified_count.get(0, 0)
            malware_rightclassified_count = rightclassified_count.get(1, 0)

            target_family_id = unknown_family_id  # 以家族ID 0和1为例
            malware_families = g.ndata['malware_family']
            # 对每个目标家族计算分类准确率

            family_test_mask = (test_mask & (malware_families == target_family_id))
            if family_test_mask.sum() > 0:
                
                family_pred = pred[family_test_mask]
                family_labels = labels[family_test_mask]
                
                true_positives = ((family_pred == 1) & (family_labels == 1)).sum().item()
                false_negatives = ((family_pred != 1) & (family_labels == 1)).sum().item()
                true_negatives = ((family_pred != 1) & (family_labels != 1)).sum().item()

                # 计算召回率
                family_recall = true_positives / (true_positives + false_negatives) if (true_positives + false_negatives) > 0 else 0

                # 计算这个家族的准确率
                total = true_positives + false_negatives + true_negatives
                family_accuracy = (true_positives + true_negatives) / total if total > 0 else 0


                print('Family ID {}: TP: {}  FN: {} TN: {}  family_recall: {:.4f} Family Acc: {:.4f}'.format(
                    target_family_id, 
                    true_positives, 
                    false_negatives, 
                    true_negatives, 
                    family_recall, 
                    family_accuracy
                ))
                # print('Family ID {}:   Family Acc: {:.3f}'.format(target_family_id, family_accuracy))

            print('Epoch {}, Loss: {:.4f}, Train Acc: {:.4f}, Test Acc: {:.4f} (Best {:.4f}), Benign Misclassified: {}, Malware Misclassified: {}, Benign Rightclassified: {}, Malware Rightclassified: {}'.format(
                e, loss, train_accurate, test_accurate, best_test_accurate, benign_misclassified_count, malware_misclassified_count, benign_rightclassified_count, malware_rightclassified_count))




if __name__ == "__main__":
    main()
